package org.example.plugin;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.*;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Signature;

/**
 *共享数据库的多租户系统实现
 * Created by yaoxiaojian
 */
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class})
})
public class MultiTenancy implements Interceptor {

    //获取tenantId的接口
    private TenantInfo tenantInfo;
    //属性参数信息
    private Properties properties;
    /**
     * 需要识别多租户字段的表名称列表
     */
    private Set<String> tableList;

    private boolean enable;

    private String tenantIdColumn;

    public MultiTenancy(){
    }

    public boolean isEnable() {
        return enable;
    }

    public void setEnable(boolean enable) {
        this.enable = enable;
    }

    public Set<String> getTableList() {
        return tableList;
    }

    public void setTableList(Set<String> tableList) {
        this.tableList = tableList;
    }

    public String getTenantIdColumn() {
        return tenantIdColumn;
    }

    public void setTenantIdColumn(String tenantIdColumn) {
        this.tenantIdColumn = tenantIdColumn;
    }

    public TenantInfo getTenantInfo() {
        return tenantInfo;
    }

    public void setTenantInfo(TenantInfo tenantInfo) {
        this.tenantInfo = tenantInfo;
    }

    public Object intercept(Invocation invocation) throws Throwable {
        if(enable) {
            mod(invocation);
        }
        return invocation.proceed();
    }

    /**
     * 更改MappedStatement为新的
     * @param invocation
     * @throws Throwable
     */
    public void mod(Invocation invocation) throws Throwable {
        MappedStatement ms = (MappedStatement) invocation
                .getArgs()[0];

        Object parameterObject = null;
        if (invocation.getArgs().length > 1) {
            parameterObject = invocation.getArgs()[1];
        }
        BoundSql boundSql = ms.getBoundSql(parameterObject);
        /**
         * 根据已有BoundSql构造新的BoundSql
         */
        BoundSql newBoundSql = new BoundSql(
                ms.getConfiguration(),
                addWhere(boundSql.getSql()),//更改后的sql
                boundSql.getParameterMappings(),
                boundSql.getParameterObject());

        /**
         * 替换MappedStatement
         */
        MappedStatement newMs=buildMappedStatement(ms,new BoundSqlSqlSource(newBoundSql));
        invocation.getArgs()[0] = newMs;
    }

    /**
     * 根据已有MappedStatement构造新的MappedStatement
     */
    private MappedStatement buildMappedStatement(MappedStatement ms, SqlSource sqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), sqlSource, ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        if (ms.getKeyProperties() != null && ms.getKeyProperties().length != 0) {
            StringBuilder keyProperties = new StringBuilder();
            for (String keyProperty : ms.getKeyProperties()) {
                keyProperties.append(keyProperty).append(",");
            }
            keyProperties.delete(keyProperties.length() - 1, keyProperties.length());
            builder.keyProperty(keyProperties.toString());
        }
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.resultSetType(ms.getResultSetType());
        builder.cache(ms.getCache());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }


    /**
     * 添加租户id条件
     * @param sql
     * @return
     * @throws JSQLParserException
     */
    private String addWhere(String sql) throws JSQLParserException {
        Statement stmt = CCJSqlParserUtil.parse(sql);
        if (stmt instanceof Insert) {
            //获得Update对象
            Insert insert = (Insert) stmt;
            Table table = insert.getTable();
            if(tableList.contains(table.getName())) {
                insert.getColumns().add(new Column(tenantIdColumn));
                ((ExpressionList) insert.getItemsList()).getExpressions().add(new StringValue("," + tenantInfo.getTenantId() + ","));
            }
            return insert.toString();
        }

        if (stmt instanceof Update) {
            Map<String,String> map = new HashMap<>();
            //获得Update对象
            Update updateStatement = (Update) stmt;
            Table table = updateStatement.getTable();
            if(table.getAlias() != null){
                map.put(table.getName(),table.getAlias().getName());
            }else{
                map.put(table.getName(),null);
            }
            List<Join> joins = updateStatement.getStartJoins();
            if(joins != null && joins.size()>0) {
                for (Join join : joins) {
                    Table table2 = (Table) join.getRightItem();
                    if (table2.getAlias() != null) {
                        map.put(table2.getName(), table2.getAlias().getName());
                    }else{
                        map.put(table.getName(),null);
                    }
                }
            }
            //获得where条件表达式
            if ( updateStatement.getWhere() instanceof BinaryExpression) {
                for (String table1 : map.keySet()) {
                    if (tableList.contains(table1)) {
                        EqualsTo equalsTo = new EqualsTo();
                        equalsTo.setLeftExpression(new Column(map.get(table1) == null ? tenantIdColumn : map.get(table1)+ "." + tenantIdColumn));
                        equalsTo.setRightExpression(new StringValue(tenantInfo.getTenantId()));
                        AndExpression andExpression = new AndExpression(equalsTo, updateStatement.getWhere());
                        updateStatement.setWhere(andExpression);
                    }
                }
            }
            return updateStatement.toString();
        }

        if (stmt instanceof Delete) {
            Map<String,String> map = new HashMap<>();
            //获得Update对象
            Delete deleteStatement = (Delete) stmt;
            //获得where条件表达式
            Table table = deleteStatement.getTable();
            if(table.getAlias() != null){
                map.put(table.getName(),table.getAlias().getName());
            }
            List<Join> joins = deleteStatement.getJoins();
            if(joins != null && joins.size()>0) {
                for (Join join : joins) {
                    Table table2 = (Table) join.getRightItem();
                    if (table2.getAlias() != null) {
                        map.put(table2.getName(), table2.getAlias().getName());
                    }
                }
            }
            if (deleteStatement.getWhere() instanceof BinaryExpression) {
                for (String table1 : map.keySet()) {
                    if (tableList.contains(table1)) {
                        EqualsTo equalsTo = new EqualsTo();
                        equalsTo.setLeftExpression(new Column(map.get(table1) == null ? table1 : map.get(table1)+ "." + tenantIdColumn));
                        equalsTo.setRightExpression(new StringValue(tenantInfo.getTenantId()));
                        AndExpression andExpression = new AndExpression(equalsTo, deleteStatement.getWhere());
                        deleteStatement.setWhere(andExpression);
                    }
                }
            }
            return deleteStatement.toString();
        }

        if (stmt instanceof Select) {
            Select select = (Select) stmt;
            processSelectBody(select.getSelectBody());
            return stmt.toString();
        }

        throw new RuntimeException("非法sql语句,请检查"+sql);
    }


    /**
     * 处理条件
     * @param expression
     * @param table
     * @return
     */
    public Expression builderExpression(Expression expression, Table table) {
        Expression tenantExpression = null;
        String[] tenantIds = {tenantInfo.getTenantId()};
        if(tableList.contains(table.getName())) {
            //当传入table时,字段前加上别名或者table名
            //别名优先使用
            StringBuilder tenantIdColumnName = new StringBuilder();
            if (table != null) {
                tenantIdColumnName.append(table.getAlias() != null ? table.getAlias().getName() : table.getName());
                tenantIdColumnName.append(".");
            }
            tenantIdColumnName.append(tenantIdColumn);
            //生成字段名
            Column tenantColumn = new Column(tenantIdColumnName.toString());

            if (tenantIds.length == 1) {
                EqualsTo equalsTo = new EqualsTo();
                tenantExpression = equalsTo;
                equalsTo.setLeftExpression(tenantColumn);
                equalsTo.setRightExpression(new StringValue("'" + tenantIds[0] + "'"));
            } else {
                //多租户身份
                InExpression inExpression = new InExpression();
                tenantExpression = inExpression;
                inExpression.setLeftExpression(tenantColumn);
                List<Expression> valueList = new ArrayList<>();
                for (String tid : tenantIds) {
                    valueList.add(new StringValue("'" + tid + "'"));
                }
                inExpression.setRightItemsList(new ExpressionList(valueList));
            }
        }
        //加入判断防止条件为空时生成 "and null" 导致查询结果为空
        if (expression == null) {
            return tenantExpression;
        } else {
            if (expression instanceof BinaryExpression) {
                BinaryExpression binaryExpression = (BinaryExpression) expression;
                getExpression(binaryExpression);
            }
            if(tenantExpression==null){
                return expression;
            }
            return new AndExpression(tenantExpression, expression);
        }

    }

    private void getExpression(BinaryExpression binaryExpression){
        if (binaryExpression.getRightExpression() instanceof FromItem) {
            processFromItem((FromItem) binaryExpression.getRightExpression());
        }
        Expression leftExpression = binaryExpression.getLeftExpression();
        if (leftExpression instanceof BinaryExpression) {
            BinaryExpression binaryExpression1 = (BinaryExpression) leftExpression;
            getExpression(binaryExpression1);

        }
        if(leftExpression instanceof  InExpression){
            InExpression inExpression = (InExpression) leftExpression;
            ItemsList rightItemsList = inExpression.getRightItemsList();
            processFromItem((FromItem) rightItemsList);
        }
    }

    /**
     * 处理子查询等
     *
     * @param fromItem
     */
    public void processFromItem(FromItem fromItem) {
        if (fromItem instanceof SubJoin) {
            SubJoin subJoin = (SubJoin) fromItem;
            if (subJoin.getJoinList() != null) {
                for (Join join : subJoin.getJoinList()) {
                    processJoin(join);
                    processFromItem(join.getRightItem());
                }
            }
            if (subJoin.getLeft() != null) {
                processFromItem(subJoin.getLeft());
            }
        } else if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            if (subSelect.getSelectBody() != null) {
                processSelectBody(subSelect.getSelectBody());
            }
        } else if (fromItem instanceof ValuesList) {

        } else if (fromItem instanceof LateralSubSelect) {
            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
            if (lateralSubSelect.getSubSelect() != null) {
                SubSelect subSelect = lateralSubSelect.getSubSelect();
                if (subSelect.getSelectBody() != null) {
                    processSelectBody(subSelect.getSelectBody());
                }
            }
        }
    }

    /**
     * 处理SelectBody
     */
    public void processSelectBody(SelectBody selectBody) {
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            if (withItem.getSelectBody() != null) {
                processSelectBody(withItem.getSelectBody());
            }
        } else {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (operationList.getSelects() != null && operationList.getSelects().size() > 0) {
                List<SelectBody> plainSelects = operationList.getSelects();
                for (SelectBody plainSelect : plainSelects) {
                    processSelectBody(plainSelect);
                }
            }
        }
    }

    /**
     * 处理PlainSelect
     */

    public void processPlainSelect(PlainSelect plainSelect) {
        processPlainSelect(plainSelect, false);
    }

    /**
     * 处理PlainSelect
     *
     * @param plainSelect
     * @param addColumn   是否添加租户列,insert into select语句中需要
     */

    public void processPlainSelect(PlainSelect plainSelect, boolean addColumn) {
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table fromTable = (Table) fromItem;
            plainSelect.setWhere(builderExpression(plainSelect.getWhere(), fromTable));
            if (addColumn)
                plainSelect.getSelectItems().add(new SelectExpressionItem(new Column("'" + this.tenantInfo.getTenantId() + "'")));

        } else {
            processFromItem(fromItem);
        }
        List<Join> joins = plainSelect.getJoins();
        if (joins != null && joins.size() > 0) {
            for (Join join : joins) {
                if(join.toString().contains("on") || join.toString().contains("ON")) {
                    processJoin(join);
                    processFromItem(join.getRightItem());
                }else {
                    Table table2 = (Table) join.getRightItem();
                    if (tableList.contains(table2.getName())) {
                        EqualsTo equalsTo = new EqualsTo();
                        equalsTo.setLeftExpression(new Column(table2.getAlias() == null ? table2.getName() : table2.getAlias().getName() + '.' + tenantIdColumn));
                        equalsTo.setRightExpression(new StringValue(tenantInfo.getTenantId()));
                        AndExpression andExpression = new AndExpression(equalsTo, plainSelect.getWhere());
                        plainSelect.setWhere(andExpression);
                    }
                }
            }
        }
    }

    /**
     * 处理联接语句
     *
     * @param join
     */
    public void processJoin(Join join) {
        if (join.getRightItem() instanceof Table) {
            Table fromTable = (Table) join.getRightItem();
                join.setOnExpression(builderExpression(join.getOnExpression(), fromTable));
        }
    }


    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    /**
     * 获取配置信息
     * {dialect=mysql, tenantInfo=org.xue.test.TenantInfoImpl, tenantIdColumn=tenant_id}
     * @param properties
     */
    public void setProperties(Properties properties) {
        this.properties=properties;
        try {
            this.enable = Boolean.parseBoolean(this.properties.getProperty("enable"));
            Class onwClass=Class.forName(this.properties.getProperty("tenantInfo"));
            this.tenantInfo=(TenantInfo)onwClass.newInstance();
        } catch (ClassNotFoundException e) {
            throw new RuntimeException(e);
        } catch (InstantiationException e) {
            throw new RuntimeException(e);
        } catch (IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     *
     */
    public static class BoundSqlSqlSource implements SqlSource {
        BoundSql boundSql;

        public BoundSqlSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }
    }
}
